'''
该代码由AI生成，需要匹配实际软硬件环境才能运行
'''

import numpy as np
from Xlib import display, X
import cv2
import time

class X11ScreenCapture:
    def __init__(self):
        self.disp = display.Display()
        self.screen = self.disp.screen()
        self.root = self.screen.root
        self.width = self.screen.width_in_pixels
        self.height = self.screen.height_in_pixels
        self.depth = self.screen.root_depth
        self.format = X.ZPixmap  # 强制使用32位格式避免24位对齐问题
        self.bytes_per_pixel = 4  # 强制使用4字节（32位）处理，避开24位的对齐坑
        self.theoretical_row_bytes = self.width * self.bytes_per_pixel

    def capture_frame(self):
        try:
            raw = self.root.get_image(
                x=0, y=0,
                width=self.width, height=self.height,
                format=self.format,
                plane_mask=0xFFFFFFFF
            )
            
            # 计算实际行字节数（兼容不同版本）
            try:
                actual_row_bytes = raw.bytes_per_line  # 优先使用行字节数属性
            except AttributeError:
                total_bytes = len(raw.data)
                actual_row_bytes = total_bytes // self.height
                print(f"使用计算的行字节数: {actual_row_bytes}")

            # 提取有效数据（去掉填充）
            data = np.frombuffer(raw.data, dtype=np.uint8)
            frame = []
            for y in range(self.height):
                start = y * actual_row_bytes
                end = start + self.theoretical_row_bytes
                row = data[start:end]
                # 确保行长度正确
                if len(row) < self.theoretical_row_bytes:
                    row = np.pad(row, (0, self.theoretical_row_bytes - len(row)), mode='constant')
                frame.append(row)
            
            # 重塑为32位图像（BGRA格式）
            frame = np.concatenate(frame).reshape((self.height, self.width, 4))
            
            # 关键修复：32位转24位BGR（OpenCV原生格式）
            bgr_frame = cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)
            
            return bgr_frame
        except Exception as e:
            print(f"捕获帧失败: {e}")
            return None

    def close(self):
        self.disp.close()

def screen_video_stream():
    capture = X11ScreenCapture()
    print(f"开始捕获视频流（分辨率：{capture.width}x{capture.height}，强制32位格式处理）")
    
    try:
        while True:
            frame = capture.capture_frame()
            if frame is not None:
                cv2.imshow("Screen Stream (Press 'q' to exit)", frame)
            
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
            time.sleep(0.01)
            
    finally:
        capture.close()
        cv2.destroyAllWindows()
        print("视频流已停止")

if __name__ == "__main__":
    screen_video_stream()